import socket
import time
import struct
from PIL import Image
from io import BytesIO
import random

NIL = 0
ACK = 1
NAK = 2

VER = "GBN"


class GBNServer:
    def __init__(self, addr: tuple[str, int], errorsize: int) -> None:
        self.serverSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.serverSocket.bind(addr)
        self.data = bytearray()
        self.tmpdata = None
        self.target = None

        # 0 -> waiting for 0 from below
        # 1 -> waiting for 1 from below
        self.state = 0
        self.errorsize = errorsize
        pass

    def recvpkt(self) -> bytes:
        data, target = self.serverSocket.recvfrom(1024)
        self.target = target
        return data

    def makepkt(self, rawdata: bytes, confirmation: int, sequence: int):
        # gbn 报文格式:
        # ack/nak 1byte
        # sequence 2byte
        # length 2byte
        # checksum 2byte
        pkt = struct.pack(
            "!BHHH%ds" % len(rawdata), confirmation, sequence, len(rawdata), 0, rawdata
        )
        csum = self.checksum(pkt)
        return struct.pack(
            "!BHHH%ds" % len(rawdata),
            confirmation,
            sequence,
            len(rawdata),
            csum,
            rawdata,
        )

    def sendpkt(self, pkt: bytes):
        self.serverSocket.sendto(pkt, self.target)

    def checksum(self, data: bytes) -> int:
        csum = 0
        for i in range(1, len(data), 2):
            csum += int.from_bytes(data[i - 1 : i + 1], "big")
        if len(data) % 2 == 1:
            csum += data[-1] << 8 & 0xFF00
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = ~csum & 0xFFFF
        # return csum
        return csum >> 8 | ((csum << 8) & 0xFF00)

    def valid(self, pkt: bytes) -> bool:
        # print('[dbg] ', self.checksum(pkt))
        return self.checksum(pkt) == 0

    def extract(self, pkt: bytes) -> tuple[int, int, int, int, bytes]:
        confirmation, sequence, length, csum = struct.unpack("!BHHH", pkt[:7])
        data = struct.unpack_from("!%ds" % length, pkt, 7)[0]
        return confirmation, sequence, length, csum, data

    def biterror(self, pkt: bytes):
        if random.randint(1, self.errorsize) == 1:
            pos = random.randint(0, len(pkt) - 1)
            off = random.randint(0, 7)
            npkt = bytearray(pkt)
            npkt[pos] ^= 1 << off
            return bytes(npkt)
        return pkt

    def run(self):
        checkedseq = 0
        while True:
            pkt2snd = self.biterror(self.makepkt(b"", ACK, checkedseq))
            pktrcvd = self.recvpkt()
            if self.valid(pktrcvd):
                confirmation, sequence, length, csum, partdata = self.extract(pktrcvd)
                if length == 0:
                    checkedseq += 1
                    print(f"responsing ack with seq {checkedseq}")
                    self.sendpkt(self.biterror(self.makepkt(b"", ACK, checkedseq)))
                    print("received all data.")
                    break
                if sequence == checkedseq + 1:
                    print(f"packet with seq {sequence} accepted.")
                    self.data += partdata
                    checkedseq += 1
                    pkt2snd = self.biterror(self.makepkt(b"", ACK, checkedseq))
                else:
                    print(f"packet with seq {sequence} rejected.")
            print(f"responsing ack with seq {checkedseq}")
            self.sendpkt(pkt2snd)


s = GBNServer(addr=("127.0.0.1", 65432), errorsize=100)
s.run()

Image.open(BytesIO(s.data)).show()

open(VER + "-server.png", "wb").write(s.data)

input("Press any key to exit...")
